import pandas as pd
import matplotlib.pyplot as plt
import os
import numpy as np

def draw_program_startup_plot(csv_file_path, output_dir):
    """
    Draw a line plot showing the relationship between total_size and avg_total_real_comm_time
    with different lines for different iterations.
    
    Args:
        csv_file_path (str): Path to the CSV file
        output_dir (str): Directory to save the output plot
    """
    # Read the CSV file
    df = pd.read_csv(csv_file_path)
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Create the plot
    plt.figure(figsize=(12, 8))
    
    # Get unique iterations and sort them
    iterations = sorted(df['iteration'].unique())
    
    # Define colors for different iterations
    colors = plt.cm.tab10(np.linspace(0, 1, len(iterations)))
    
    # Plot lines for each iteration
    for i, iteration in enumerate(iterations):
        iteration_data = df[df['iteration'] == iteration]
        
        # Sort by total_size for proper line connection
        iteration_data = iteration_data.sort_values('total_size')
        
        plt.plot(iteration_data['total_size'], 
                iteration_data['avg_total_real_comm_time'],
                marker='o', 
                linewidth=2, 
                markersize=6,
                color=colors[i],
                label=f'Iteration {iteration}')
    
    # Customize the plot
    plt.xlabel('total_size (B)', fontsize=12)
    plt.ylabel('Average program_startup_cost Time (us)', fontsize=12)
    plt.title('Program Startup Cost for Different Iterations', fontsize=14)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, alpha=0.3)
    
    # Adjust layout to prevent legend cutoff
    plt.tight_layout()
    
    # Save the plot
    output_file = os.path.join(output_dir, 'program_startup_cost_by_iteration.png')
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    
    # Also save as PDF for better quality
    output_file_pdf = os.path.join(output_dir, 'program_startup_cost_by_iteration.pdf')
    plt.savefig(output_file_pdf, bbox_inches='tight')
    
    print(f"Plot saved to: {output_file}")
    print(f"Plot saved to: {output_file_pdf}")
    
    # Show the plot
    plt.show()
    
    # Print some statistics
    print("\nData Summary:")
    print(f"Total data points: {len(df)}")
    print(f"Iterations: {iterations}")
    print(f"Total size range: {df['total_size'].min():.1f} - {df['total_size'].max():.1f}")
    print(f"Communication time range: {df['avg_total_real_comm_time'].min():.2f} - {df['avg_total_real_comm_time'].max():.2f}")

def main():
    # Define base directory and target directory
    base_dir = r"F:\PostGraduate\Point-to-Point-DATA\deal-data-code\C-lop-Prediction\analysis_for_sendrecv\analysis_4node_all_iteration"
    target_dir = r"F:\PostGraduate\Point-to-Point-DATA\deal-data-code\C-lop-Prediction\analysis_for_sendrecv\analysis_4node_all_iteration\plots"
    
    # CSV file path
    csv_file = os.path.join(base_dir, "program_startup_cost.csv")
    
    # Check if CSV file exists
    if not os.path.exists(csv_file):
        print(f"Error: CSV file not found at {csv_file}")
        return
    
    # Draw the plot
    draw_program_startup_plot(csv_file, target_dir)

if __name__ == "__main__":
    main()